# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
""" Produce the dataset for mnist. """

import os
from urllib.error import URLError

from mindspore.dataset.vision import Inter
from mindspore.common import dtype as mstype
import mindspore.dataset.vision.c_transforms as C
import mindspore.dataset.transforms.c_transforms as C2

from mindvision.common.dataset.dataloader import DataLoader
from mindvision.classification.dataset.meta_dataset import Meta


class MnistDataLoader(DataLoader, metaclass=Meta):
    """Mnist Dataset"""
    mirrors = [
        'http://yann.lecun.com/exdb/mnist/',
        'https://ossci-datasets.s3.amazonaws.com/mnist/',
    ]

    resources = [
        ("train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"),
        ("train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"),
        ("t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"),
        ("t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c")
    ]

    index2label = {0: 'zero', 1: 'one', 2: 'two', 3: 'three', 4: 'four', 5: 'five', 6: 'six', 7: 'seven', 8: 'eight',
                   9: 'nine'}

    def __init__(self, dataset=None, train=True, transform=None, target_transform=None, batch_size=32, repeat_num=1,
                 num_parallel_workers=None, download=False):

        super(MnistDataLoader, self).__init__(dataset=dataset, train=train, transform=transform,
                                              target_transform=target_transform, batch_size=batch_size,
                                              repeat_num=repeat_num, num_parallel_workers=num_parallel_workers,
                                              download=download)

        if download:
            self._download()

    def _download(self):
        """Download the MNIST data if it doesn't exist already"""
        if self._check_exists():
            return

        os.makedirs(self.raw_folder, exist_ok=True)

        # download files
        for filename, md5 in self.resources:
            for mirror in self.mirrors:
                url = "{}{}".format(mirror, filename)
                try:
                    print("Downloading {}".format(url))
                    self._download_url(
                        url, root=self.raw_folder,
                        filename=filename,
                        md5=md5
                    )
                except URLError as error:
                    print("Failed to download (trying next):\n{}".format(error))
                    continue
                finally:
                    print()
                break
            else:
                raise RuntimeError("Error downloading {}".format(filename))

    def _default_transform(self):
        """default transform"""
        resize_height = 32
        resize_width = 32
        rescale = 1.0 / 255.0
        shift = 0.0
        rescale_nml = 1 / 0.3081
        shift_nml = -1 * 0.1307 / 0.3081

        # define map operations
        trans = [
            C.Resize((resize_height, resize_width), interpolation=Inter.LINEAR),
            C.Rescale(rescale, shift),
            C.Rescale(rescale_nml, shift_nml),
            C.HWC2CHW(),
        ]
        return trans

    def _transforms(self):
        """transforms"""
        assert self.dataset, "dataset is None"
        trans = self.transform if self.transform else self._default_transform()
        self.dataset = self.dataset.map(operations=trans, input_columns="image",
                                        num_parallel_workers=self.num_parallel_workers)
        type_cast_op = self.target_transform if self.target_transform else C2.TypeCast(mstype.int32)
        self.dataset = self.dataset.map(operations=type_cast_op,
                                        input_columns="label",
                                        num_parallel_workers=self.num_parallel_workers)

    def _pipeline(self):
        """pipeline"""
        self._transforms()
        self.dataset = self.dataset.batch(self.batch_size, drop_remainder=True)
        self.dataset = self.dataset.repeat(self.repeat_num)
